import sys
import os

# 添加 project/ 目录到 sys.path
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

import argparse

import cv2
import numpy as np
from tqdm import tqdm
import sys

sys.path.append('../')  # 确保这个路径是正确的
from cyw_devkit.core.dataset import CywDataset
from cyw_devkit.core.visualizer import draw_boxes2img, draw_pts2img
from tictoc import TicToc
import multiprocessing


def main(args):
    # 统计耗时
    cost = TicToc("数据可视化")
    assert os.path.exists(args.data_path)
    frames = os.listdir(args.data_path)
    frames.sort(key=lambda x: x)
    files = []
    for dir in frames:
        if dir[:2] == '__':  # '__'
            dir = os.path.join(args.data_path, dir)
            if os.path.isdir(dir):
                files.append(dir)

    process_size = len(files)
    # for idx in range(process_size):
    #     main_worker(files[0], args)

    manager = multiprocessing.Manager()
    if process_size > 1:
        pool = multiprocessing.Pool(process_size)
        counter_list = manager.list()
        for idx in range(process_size):
            pool.apply_async(main_worker, args=(files[idx], args))
        pool.close()
        pool.join()
    else:
        main_worker(files[0], args)

    print("---------------------------------------------------------")
    print("处理完成: {}".format(files))
    cost.toc()
    print("---------------------------------------------------------")


def main_worker(dataset_path, args):
    cyw = CywDataset(dataset_path=dataset_path, lidar_dim=4, radar_dim=7, normal_img=False)
    # output_file = os.path.join(args.out_path,os.path.basename(dataset_path)+'.mp4')
    # fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # 设置视频编码格式（这里使用 MP4V 编码）
    # out = cv2.VideoWriter(output_file, fourcc, 30, (1080, 720))

    for idx, frame_data in enumerate(tqdm(cyw)):
        # if idx < 100:
        #     continue
        img_dict = dict()
        for camera_id in frame_data['camera_dict'].keys():
            lidar = frame_data['lidar']
            image = frame_data['camera_dict'][camera_id]
            if args.sync_time:  # 是否考虑传感器之间的时间差
                camera2base = cyw.data_dict['super_transform'].get_tf(source_point=image.header.frame_id,
                                                                      source_time=image.header.stamp,
                                                                      target_point=lidar.header.frame_id,
                                                                      target_time=lidar.header.stamp)
            else:
                camera2base = cyw.data_dict['super_transform'].get_tf_from_space(source_point=camera_id,
                                                                                 target_point='base_link')

            camera2img = cyw.get_camera_mat(camera_id)

            image_new = draw_pts2img(points=lidar.data[::, ], image=np.array(image.data),
                                     tfs=[np.linalg.inv(camera2base), camera2img])  # 跳一下点，不然速度太慢

            label = frame_data['label']
            if label is not None:
                image_new = draw_boxes2img(boxes=label.data.tensors, image=image_new,
                                           tfs=[np.linalg.inv(camera2base), camera2img])

            img_dict[camera_id] = image_new

        img_height = image_new.shape[0]
        img_width = image_new.shape[1]
        # 创建一个新的大图像，用于放置九宫格布局
        grid_img = np.zeros((2 * img_height, 3 * img_width, 3), dtype=np.uint8)
        # 定义每个小图像的位置（索引从0开始，位置从1开始）
        positions = {
            'camera75': (0, 1),  # 前
            'camera81': (0.5, 0),  # 左
            'camera80': (0.5, 2),  # 右
            'camera77': (1, 1)  # 后
        }
        # 将每个图像放置到对应的九宫格位置，其他位置保持黑色
        for id, (row, col) in positions.items():
            if id in img_dict.keys():
                grid_img[int(row * img_height):int((row + 1) * img_height),
                int(col * img_width):int((col + 1) * img_width)] = img_dict[id]

        #         # 保存合并后的图像
        # cv2.imwrite('combined_image.jpg', grid_img)
        #
        # # 显示合并后的图像（可选）
        # cv2.imshow('Combined Image', grid_img)
        # cv2.waitKey(0)
        # cv2.destroyAllWindows()
        grid_img = cv2.resize(grid_img, (1080, 720))
        cv2.imwrite(os.path.join(args.out_path, f'{os.path.basename(dataset_path)}_{idx}.png'), grid_img)
        # just_look_array(img=grid_img, title=str(idx))
    #     out.write(grid_img)
    # out.release()


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Configuration Parameters')
    parser.add_argument('--data-path', default="/home/adt/bags/work_space/datasets",
                        help='your data root')
    parser.add_argument('--out-path', default="/home/adt/bags/work_space/dataview",
                        help='')
    parser.add_argument('--sync-time', action='store_true',
                        help='')
    args = parser.parse_args()

    main(args)
